#!/usr/bin/env python
# encoding: utf-8

import argparse
import torch
from sklearn.model_selection import train_test_split

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_data_list_path', type=str)
    parser.add_argument('--output_train_list_path', type=str)
    parser.add_argument('--output_val_list_path', type=str)
    parser.add_argument('--val_size', type=float, default=0.05)
    args = parser.parse_args()

    audioSequences, Speakers = torch.load(args.input_data_list_path)
    train_seq, val_seq = train_test_split(audioSequences, test_size=0.05)

    torch.save((train_seq, Speakers), args.output_train_list_path)
    print(f'Saved data list file at {args.output_train_list_path}')
    torch.save((val_seq, Speakers), args.output_val_list_path)
    print(f'Saved data list file at {args.output_val_list_path}')


